import torch.utils.data
from . import dataset as D


def create_dataloader(dataset, phase):
    '''create dataloader '''
    assert phase in ["train", "test"], "phase must be \"train\" or \"test\""
    if phase == 'train':
        return torch.utils.data.DataLoader(
            dataset,
            batch_size=10,
            shuffle=True,
            num_workers=16,
            drop_last=True,
            pin_memory=True)
    else:
        return torch.utils.data.DataLoader(
            dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=True)


def create_dataset(mode, phase, exp):
    '''create dataset'''
    dataset = None
    assert mode in ["LR", "LRHR"], "mode must be \"LRHR\" or \"LR\""
    if mode == "LR":
        dataset = D.LRDataset(exp)
    elif mode == "LRHR":
        dataset = D.LRHRDataset(phase, exp)

    if dataset:
        print("Datasets have been created.")

    return dataset
